from tensorflow.keras import Model
from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, Dense, Dropout
import numpy as np

"""
tf.keras.layers.Conv2D
parameters:
- input_shape: 64*64 color pics (64, 64, 3)
- filters: conv number
- kernel_size: kernel_size(3, 3)
- strides: strides=(2, 2)
- padding: 'VALID', 'SAME'(hights&weights of feature map = ceil(h&w of pic / strides)
- activation: 'sigmoid', 'tanh', 'relu', 'elu', 'selu', 'swish', 'softmax'
- input_shape: only for the first layer , (height, weight, channel)
- name: layer names assignment

tf.keras.layers.Maxpool2D
parameters:
- pool_size: the same as kernel_size in Conv2D
- strides
- name

tf.keras.layers.Dense
parameters:
- units: neuron number
- activation
- name

tf.keras.layers.Flatten
[[1, 2, 3],
[4, 5, 6],   ---------->  [1, 2, 3, 4, 5, 6, 7, 8, 9]
[7, 8, 9]]
"""


class ConvModel(Model):

    def __init__(self):
        """
        model structure
        -----------------------------------------------------
        \ names \ filters \ kernel_size\activation\input_size
        -----------------------------------------------------
        \ conv_1_1 \ 64   \ 3          \ relu     \160*160*3
        \ maxpool  \      \ 2          \          \160*160*64
        \ con_2_1  \ 128  \ 3          \ relu     \80*80*64
        \ maxpool  \      \ 2          \          \80*80*128
        \ con_3_1  \ 256  \ 3          \ relu     \40*40*128
        \ con_3_2  \ 256  \ 3          \ relu     \40*40*256
        \ maxpool  \      \ 2          \          \40*40*256
        \ con_4_1  \ 512  \ 3          \ relu     \20*20*256
        \ con_4_2  \ 512  \ 3          \ relu     \20*20*512
        \ maxpool  \      \ 2          \          \20*20*512
        \ con_5_1  \ 512  \ 3          \ relu     \10*10*512
        \ con_5_2  \ 512  \ 3          \ relu     \10*10*512
        \ maxpool  \      \ 2          \          \10*10*512
        \ Fatten   \      \            \          \5*5*512
        \ Dense_1  \ 4096 \            \ relu     \12800
        \ dropout
        \ Dense_2  \ 4096 \            \ relu     \4096
        \ dropout
        \ Dense_3  \ 10   \            \ logit    \4096
        ------------------------------------------------------
        """

        super(ConvModel, self).__init__()
        np.random.seed(7)

        self.conv_1_1 = Conv2D(input_shape=(160, 160, 3), filters=64, kernel_size=3,
                               padding='same', activation='relu', kernel_initializer="he_normal", name='conv_1_1')

        self.max_pool_1 = MaxPooling2D(pool_size=2, name='max_pool_1')

        self.conv_2_1 = Conv2D(filters=128, kernel_size=3, padding='same',
                               activation='relu', kernel_initializer="he_normal", name='conv_2_1')

        self.max_pool_2 = MaxPooling2D(pool_size=2, name='max_pool_2')

        self.conv_3_1 = Conv2D(filters=256, kernel_size=3, padding='same',
                               activation='relu', kernel_initializer="he_normal", name='conv_3_1')

        self.conv_3_2 = Conv2D(filters=256, kernel_size=3, padding='same',
                               activation='relu', kernel_initializer="he_normal", name='conv_3_2')

        self.max_pool_3 = MaxPooling2D(pool_size=2, name='max_pool_3')

        self.conv_4_1 = Conv2D(filters=512, kernel_size=3, padding='same',
                               activation='relu', kernel_initializer="he_normal", name='conv_4_1')

        self.conv_4_2 = Conv2D(filters=512, kernel_size=3, padding='same',
                               activation='relu', kernel_initializer="he_normal", name='conv_4_2')

        self.max_pool_4 = MaxPooling2D(pool_size=2, name='max_pool_4')

        self.conv_5_1 = Conv2D(filters=512, kernel_size=3, padding='same',
                               activation='relu', kernel_initializer="he_normal", name='conv_5_1')

        self.conv_5_2 = Conv2D(filters=512, kernel_size=3, padding='same',
                               activation='relu', kernel_initializer="he_normal", name='conv_5_2')

        self.max_pool_5 = MaxPooling2D(pool_size=2, name='max_pool_5')

        self.flatten = Flatten(name='flatten')

        self.dense_1 = Dense(units=4096, activation="relu", kernel_initializer="he_normal", name='dense_1')

        self.dropout = Dropout(0.5)

        self.dense_2 = Dense(units=4096, activation="relu", kernel_initializer="he_normal", name='dense_2')

        self.dropout = Dropout(0.5)

        self.dense_3 = Dense(units=10, activation="softmax", kernel_initializer="he_normal", name='logit')

    def call(self, inputs, training=None):
        x = self.conv_1_1(inputs)
        x = self.max_pool_1(x)
        x = self.conv_2_1(x)
        x = self.max_pool_2(x)
        x = self.conv_3_1(x)
        x = self.conv_3_2(x)
        x = self.max_pool_3(x)
        x = self.conv_4_1(x)
        x = self.conv_4_2(x)
        x = self.max_pool_4(x)
        x = self.conv_5_1(x)
        x = self.conv_5_2(x)
        x = self.max_pool_5(x)
        x = self.flatten(x)
        x = self.dense_1(x)
        x = self.dropout(x)
        x = self.dense_2(x)
        x = self.dropout(x)
        x = self.dense_3(x)
        return x



